import requests
import math
import csv
import time

# --- 1. 配置你的请求参数 ---
# !!! 重要：请将下面的Key替换为您自己在高德开放平台申请的【Web服务】Key !!!
api_key = "dc855d1035cae707d1fa0ed4ec7009a6"
city = "福州"
offset = 25

def fetch_all_pois_weighted(keywords, poi_type_code, output_filename, assigned_weight):
    """
    一个通用的函数，用于获取指定类型的所有POI数据，并按指定格式保存。
    
    :param keywords: 搜索关键词
    :param poi_type_code: POI的分类代码
    :param output_filename: 输出的CSV文件名
    :param assigned_weight: 为此类别的所有条目指定的固定权重值
    """
    print(f"\n{'='*50}")
    print(f"▶️ 开始查找: {keywords} (权重: {assigned_weight})")
    print(f"{'='*50}")

    all_pois = []
    page_num = 1

    # --- 首次请求获取总数 ---
    url = "https://restapi.amap.com/v3/place/text"
    params = {
        'key': api_key,
        'keywords': keywords,
        'types': poi_type_code,
        'city': city,
        'citylimit': 'true',
        'offset': offset,
        'page': page_num
    }
    
    try:
        response = requests.get(url, params=params, timeout=10)
        response.raise_for_status()
        first_page_data = response.json()
    except requests.exceptions.RequestException as e:
        print(f"❌ 首次请求失败: {e}")
        return

    if first_page_data and first_page_data['status'] == '1':
        total_count = int(first_page_data.get('count', 0))
        if total_count == 0:
            print(f"✅ 查询成功，但未找到任何相关的 '{keywords}'。")
            return
        
        print(f"✅ 查询成功！共找到 {total_count} 个 '{keywords}'。")
        all_pois.extend(first_page_data.get('pois', []))
        
        # --- 计算总页数并循环获取所有数据 ---
        total_pages = math.ceil(total_count / offset)
        print(f"   每页最多显示 {offset} 条，共需要请求 {total_pages} 页。")

        if total_pages > 1:
            for page in range(2, total_pages + 1):
                print(f"   ...正在获取第 {page}/{total_pages} 页数据...")
                params['page'] = page
                try:
                    time.sleep(0.1)
                    response = requests.get(url, params=params, timeout=10)
                    response.raise_for_status()
                    page_data = response.json()
                    if page_data and page_data.get('pois'):
                        all_pois.extend(page_data['pois'])
                except requests.exceptions.RequestException as e:
                    print(f"   ❌ 获取第 {page} 页数据时失败: {e}")

        # --- 将所有数据按指定格式写入CSV文件 ---
        if not all_pois:
            print("   未能获取到任何POI信息，无法生成CSV文件。")
            return
            
        print(f"\n   正在将 {len(all_pois)} 条数据写入文件: {output_filename} ...")
        try:
            with open(output_filename, mode='w', newline='', encoding='utf-8-sig') as file:
                # 使用您指定的表头
                csv_headers = ['name', 'address', 'geolocation_lng', 'geolocation_lat', 'category', 'weight']
                writer = csv.writer(file)
                writer.writerow(csv_headers)
                
                for poi in all_pois:
                    location_str = poi.get('location', ',')
                    lng, lat = (location_str.split(',') + ['N/A', 'N/A'])[:2]
                    row = [
                        poi.get('name', 'N/A'),
                        poi.get('address', 'N/A'),
                        lng,
                        lat,
                        poi.get('type', 'N/A'), # 'type' 字段对应 category
                        assigned_weight # 使用传入的固定权重
                    ]
                    writer.writerow(row)
            print(f"   🎉 成功！数据已保存到: {output_filename}")
        except IOError as e:
            print(f"   ❌ 写入文件时发生错误: {e}")
    else:
        print(f"❌ 高德API返回错误。状态码: {first_page_data.get('status')}, 错误信息: {first_page_data.get('info')}")


if __name__ == "__main__":
    if api_key == "请在这里粘贴你自己的Key":
        print("🛑 错误：请在脚本的第7行将 'api_key' 替换为您自己的高德API Key！")
    else:
        # --- 任务1: 查找写字楼, 指定权重为 1.0 ---
        fetch_all_pois_weighted(
            keywords="仓山区 写字楼", 
            poi_type_code="120200", 
            output_filename="cangshan_office_buildings_weighted.csv",
            assigned_weight=1.0
        )
        
        # --- 任务2: 查找学校, 指定权重为 1.5 ---
        fetch_all_pois_weighted(
            keywords="仓山区 学校", 
            poi_type_code="141200", 
            output_filename="cangshan_schools_weighted.csv",
            assigned_weight=1.5
        )
        
        print(f"\n{'='*50}")
        print("✅✅✅ 所有任务执行完毕！ ✅✅✅")
        print(f"{'='*50}")